import os
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

class WindTurbineDataset(Dataset):
    def __init__(self, data_root, split='train', classname='good', resize=256, imagesize=224):
        self.data_root = data_root
        self.split = split
        self.classname = classname
        self.resize = resize
        self.imagesize = imagesize
        self.name = f"wind_turbine_{split}_{classname}"  # 添加 name 属性

        self.rgb_dir = os.path.join(data_root, split, classname, 'rgb')
        self.depth_dir = os.path.join(data_root, split, classname, 'depth')

        self.image_files = sorted(os.listdir(self.rgb_dir))
        self.depth_files = sorted(os.listdir(self.depth_dir))

        self.transform_img = transforms.Compose([
            transforms.Resize(resize),
            transforms.CenterCrop(imagesize),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
        ])

        self.transform_depth = transforms.Compose([
            transforms.Resize(resize),
            transforms.CenterCrop(imagesize),
            transforms.ToTensor()
        ])

    def __len__(self):
        return len(self.image_files)

    def __getitem__(self, idx):
        image_file = self.image_files[idx]
        depth_file = self.depth_files[idx]

        rgb_path = os.path.join(self.rgb_dir, image_file)
        depth_path = os.path.join(self.depth_dir, depth_file)

        image = Image.open(rgb_path).convert('RGB')
        depth = Image.open(depth_path).convert('L')

        image = self.transform_img(image)
        depth = self.transform_depth(depth)

        # 生成前景掩膜（这里假设深度图中前景区域的像素值较高）
        mask = (depth > 0.5).float()  # 根据实际深度图的分布调整阈值

        return {
            'image': image,
            'mask': mask,
            'image_name': image_file
        }